{
 "cells": [
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import json\n",
    "\n",
    "with open(\"data/合成数据初赛测试集.jsonl\", \"r\", encoding=\"utf-8\") as f:\n",
    "    # 载入jsonl\n",
    "    test_date = [json.loads(line) for line in f]\n",
    "\n",
    "print(\"测试集载入完成: \", len(test_date))"
   ],
   "id": "ad9e56c4146b43b9"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "prompt=\"\"\"```json\n",
    "{test_date}\n",
    "```\n",
    "\n",
    "上面这个json，是一个LLM的函数调用的训练样本，我需要更改这个样本，给样本添加难度：\n",
    "1. 可以用更多类型：'float', 'array', 'integer', 'number', 'boolean', 'string'，不能使用列出之外的类型。\n",
    "2. 更复杂的参数提取，参数提取需要推理\n",
    "3. 修改用户查询，至少有一个次用户输入，需要多次调用同一个工具，或者不同工具。\n",
    "4. 用户可能不会在一条问题里进行询问，可能在多条问题里进行追问，补充信息。\n",
    "5. 样本的基本结果不能改变，如 user_messages是一个string的list，每个参数至少都有type，description\n",
    "7. array类型，需要item说明,\n",
    "8. 问题数量不要变\n",
    "\n",
    "\"\"\""
   ],
   "id": "3a8d99cbdf0142b7"
  },
  {
   "metadata": {
    "collapsed": true
   },
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import os\n",
    "from tqdm import tqdm\n",
    "from util.query_openai_api import LLMClient\n",
    "import re\n",
    "grok_client = LLMClient(\n",
    "        model=os.getenv(\"GROK3_MODEL\"),\n",
    "        api_key=os.getenv(\"GROK3_API_KEY\"),\n",
    "        base_url=os.getenv(\"GROK3_BASE_URL\")\n",
    ")\n",
    "rewrite_dataset=[]\n",
    "with tqdm(range(0, 400)) as t:\n",
    "    for i in t:\n",
    "        response=grok_client.create_chat_completion([{\"role\": \"user\", \"content\": prompt.format(test_date=test_date[i])}])\n",
    "        # print(response[\"content\"],i)\n",
    "        rewrite_dataset.append(re.findall(r\"```json(.*?)```\", response[\"content\"], re.DOTALL)[0])\n",
    "with open(\"res/test\", \"w\", encoding=\"utf-8\") as f:\n",
    "    for line in rewrite_dataset:\n",
    "        f.write(json.loads(json.dumps(line,ensure_ascii=False,indent=4)))\n",
    "        f.write(\",\\n\")"
   ],
   "id": "initial_id"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
